import cv2
import os
import sys
import pickle
import numpy as np

# 常量定义
IMAGE_DIR = "Images"  # 路牌图片库文件夹·
CAPTURE_PATH = "capture.jpg"  # 拍照保存路径
IMAGE_SIZE = (128, 128)  # 统一图片大小
FEATURE_LIB_PATH = "feature_library.pkl"  # 特征库路径

# 进度条符号（科幻风格）
PROGRESS_SYMBOLS = ['░', '▒', '▓', '█']
BAR_LENGTH = 30

# 定义颜色的HSV范围
color_ranges = {
    "red": [(np.array([0, 150, 150]), np.array([10, 255, 255])),
            (np.array([170, 150, 150]), np.array([180, 255, 255]))],
    "green": [(np.array([36, 100, 100]), np.array([86, 255, 255]))],
    "blue": [(np.array([94, 80, 2]), np.array([126, 255, 255]))],
    "yellow": [(np.array([15, 150, 150]), np.array([35, 255, 255]))]
}

def sci_fi_progress_bar(progress, total, prefix="识别进度"):
    percent = int(progress / total * 100)
    filled_len = int(BAR_LENGTH * progress // total)
    bar = ""
    for i in range(BAR_LENGTH):
        if i < filled_len:
            bar += PROGRESS_SYMBOLS[(i // 2) % len(PROGRESS_SYMBOLS)]
        else:
            bar += ' '
    sys.stdout.write(f"\r\033[96m{prefix}: |{bar}| {percent}%\033[0m")
    sys.stdout.flush()

def orb_match(img1, img2):
    orb = cv2.ORB_create(
        nfeatures=1000,
        scaleFactor=1.2,
        nlevels=8,
        edgeThreshold=31
    )
    kp1, des1 = orb.detectAndCompute(img1, None)
    kp2, des2 = orb.detectAndCompute(img2, None)
    if des1 is None or des2 is None or len(des1) < 2 or len(des2) < 2:
        return 0
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    matches = bf.match(des1, des2)
    good_matches = [m for m in matches if m.distance < 50]
    return len(good_matches)


def get_chinese_label(filename):
    if "Left" in filename:
        return "左转"
    elif "Right" in filename:
        return "右转"
    elif "Stop" in filename:
        return "停车位"
    elif "Parking" in filename:
        return "禁止停车"


def load_feature_library():
    """加载预训练的特征库"""
    if not os.path.exists(FEATURE_LIB_PATH):
        print("\n\033[91m错误：特征库文件不存在！请先运行 train_feature_library.py 训练特征库\033[0m")
        return None
    try:
        with open(FEATURE_LIB_PATH, 'rb') as f:
            return pickle.load(f)
    except Exception as e:
        print(f"\n\033[91m错误：加载特征库失败 - {str(e)}\033[0m")
        return None


def preprocess_image(img):
    """图像预处理，增强特征提取效果"""
    img_eq = cv2.equalizeHist(img)
    img_blur = cv2.GaussianBlur(img_eq, (3, 3), 0)
    return img_blur


def compare_with_library(capture_path, feature_library):
    """使用预训练的特征库进行识别"""
    if feature_library is None:
        return

    img1 = cv2.imread(capture_path, cv2.IMREAD_GRAYSCALE)
    if img1 is None:
        print("\n\033[91m错误：无法读取拍摄的图片！\033[0m")
        return

    img1 = cv2.resize(img1, IMAGE_SIZE)
    img1_processed = preprocess_image(img1)

    orb = cv2.ORB_create(
        nfeatures=2000,
        scaleFactor=1.1,
        nlevels=12,
        edgeThreshold=15,
        firstLevel=0,
        WTA_K=2,
        patchSize=31,
        fastThreshold=20
    )

    kp1, des1 = orb.detectAndCompute(img1_processed, None)
    if des1 is None or len(des1) < 10:
        print("\n\033[91m错误：拍摄的图片特征点不足！\033[0m")
        return

    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    match_scores = []
    total = len(feature_library)
    print("\n\033[95m正在进行路牌识别比对...\033[0m")

    for idx, (fname, features) in enumerate(feature_library.items(), 1):
        des2 = features['descriptors']
        kp_count = features['keypoints']
        matches = bf.match(des1, des2)

        if len(matches) > 0:
            distances = [m.distance for m in matches]
            mean_dist = np.mean(distances)
            std_dist = np.std(distances)
            good_matches = [m for m in matches if m.distance < mean_dist - 0.5 * std_dist]
            score = len(good_matches) * (1.0 - mean_dist / 100.0)
            match_scores.append((fname, score, kp_count, mean_dist, len(good_matches)))

        progress = idx / total * 100
        print(f"\r\033[96m比对进度: {progress:.1f}% ({idx}/{total})\033[0m", end='')

    print()

    if not match_scores:
        print("\n\033[91m错误：没有找到任何匹配！\033[0m")
        return

    match_scores.sort(key=lambda x: x[1], reverse=True)
    best_name, best_score, kp_count, mean_dist, good_matches = match_scores[0]

    print("\n\033[95m匹配详情：\033[0m")
    print("\n\033[96m最佳匹配结果（前5名）：\033[0m")
    for i, (fname, score, kp_count, mean_dist, good_matches) in enumerate(match_scores[:5], 1):
        label = get_chinese_label(fname)
        print(f"{i}. {label} ({fname}):")
        print(f"   - 匹配分数: {score:.2f}")
        print(f"   - 特征点数量: {kp_count}")
        print(f"   - 平均匹配距离: {mean_dist:.2f}")
        print(f"   - 好的匹配点数: {good_matches}")
        print(f"   - 匹配率: {good_matches / kp_count * 100:.1f}%")

    if good_matches < 5 or mean_dist > 50:
        print("\n\033[93m警告：匹配质量较低，结果可能不可靠\033[0m")

    label = get_chinese_label(best_name)
    print(f"\n\033[92m最终识别结果：{label}\033[0m")
    print(f"\033[92m匹配图片：{best_name}\033[0m")
    print(f"\033[92m最佳匹配分数：{best_score:.2f}\033[0m")


def road_sign_recognition():
    """路牌识别功能"""
    import time
    feature_library = load_feature_library()
    if feature_library is None:
        return

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("无法打开摄像头")
        return

    window_name = "路牌识别 - ESC退出"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    # 设置窗口置顶
    try:
        cv2.setWindowProperty(window_name, cv2.WND_PROP_TOPMOST, 1)
    except Exception:
        pass

    print("\033[94m摄像头已开启，2秒后自动拍照识别，按ESC退出\033[0m")
    start_time = time.time()
    captured = False
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        cv2.imshow(window_name, frame)
        key = cv2.waitKey(1) & 0xFF
        if key == 27:  # ESC键退出
            print("\n退出路牌识别模式")
            break
        # 2秒后自动拍照
        if not captured and (time.time() - start_time) >= 2.0:
            cv2.imwrite(CAPTURE_PATH, frame)
            print("\n\033[93m已自动拍照并保存为 capture.jpg\033[0m")
            compare_with_library(CAPTURE_PATH, feature_library)
            captured = True
            # 等待用户按ESC退出
    cap.release()
    cv2.destroyAllWindows()


def color_detection():
    """颜色识别功能"""
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("无法打开摄像头")
        return

    window_name = '颜色识别 - 按ESC退出'
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    # 设置窗口置顶
    try:
        cv2.setWindowProperty(window_name, cv2.WND_PROP_TOPMOST, 1)
    except Exception:
        pass

    print("\033[94m颜色识别模式已启动，按ESC退出\033[0m")

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

        for color, ranges in color_ranges.items():
            mask = None
            for lower, upper in ranges:
                if mask is None:
                    mask = cv2.inRange(hsv, lower, upper)
                else:
                    mask = cv2.bitwise_or(mask, cv2.inRange(hsv, lower, upper))

            # 形态学操作去噪
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))

            # 查找轮廓
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            for cnt in contours:
                area = cv2.contourArea(cnt)
                if area > 500:  # 过滤小面积
                    x, y, w, h = cv2.boundingRect(cnt)
                    cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    cv2.putText(frame, color, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)

        cv2.imshow(window_name, frame)
        if cv2.waitKey(1) & 0xFF == 27:  # 按ESC退出
            print("\n退出颜色识别模式")
            break

    cap.release()
    cv2.destroyAllWindows()


def show_menu():
    """显示主菜单"""
    print("\n" + "=" * 50)
    print("\033[95m          小车视觉识别系统\033[0m")
    print("=" * 50)
    print("\033[96m请选择功能：\033[0m")
    print("1. 路牌识别模块（按空格拍照）")
    print("2. 小球颜色识别模块")
    print("3. 退出程序")
    print("=" * 50)


def main():
    """主函数"""
    while True:
        show_menu()
        choice = input("\033[93m请输入选择 (1-3): \033[0m").strip()

        if choice == "1":
            print("\n\033[94m启动路牌识别功能...\033[0m")
            road_sign_recognition()
        elif choice == "2":
            print("\n\033[94m启动颜色识别功能...\033[0m")
            color_detection()
        elif choice == "3":
            print("\n\033[92m感谢使用，再见！\033[0m")
            break
        else:
            print("\n\033[91m无效选择，请重新输入！\033[0m")


if __name__ == "__main__":
    main() 